"""
Mesh optimization module
"""
import numpy as np
from .mesh_quality import MeshQuality

class MeshOptimizer:
    """Class for optimizing mesh quality"""
    
    def __init__(self, mesh_data):
        """
        Initialize the mesh optimizer
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
        """
        self.mesh_data = mesh_data
        self.original_mesh_data = self._copy_mesh_data(mesh_data)
        self.quality_metrics = MeshQuality.calculate_quality_metrics(mesh_data)
    
    def _copy_mesh_data(self, mesh_data):
        """Create a deep copy of mesh data"""
        copy = {}
        
        if 'points' in mesh_data:
            copy['points'] = np.copy(mesh_data['points'])
        
        if 'cells' in mesh_data:
            copy['cells'] = [cell.copy() for cell in mesh_data['cells']]
        
        # Don't copy VTK objects, they'll be recreated when needed
        
        return copy
    
    def optimize(self, method='laplacian', iterations=10, quality_threshold=2.0):
        """
        Optimize the mesh using the specified method
        
        Parameters:
        -----------
        method : str
            Optimization method ('laplacian', 'smart', or 'ai')
        iterations : int
            Number of optimization iterations
        quality_threshold : float
            Quality threshold for identifying poor elements
            
        Returns:
        --------
        dict
            Dictionary containing optimized mesh data and quality improvement statistics
        """
        if method == 'laplacian':
            return self._laplacian_smoothing(iterations, quality_threshold)
        elif method == 'smart':
            return self._smart_smoothing(iterations, quality_threshold)
        elif method == 'ai':
            return self._ai_optimization(iterations, quality_threshold)
        else:
            raise ValueError(f"Unsupported optimization method: {method}")
    
    def _laplacian_smoothing(self, iterations, quality_threshold):
        """
        Optimize mesh using Laplacian smoothing
        
        This method moves each vertex to the average position of its neighbors.
        """
        # Create a copy of the mesh data to work with
        optimized_mesh = self._copy_mesh_data(self.mesh_data)
        points = optimized_mesh['points']
        cells = optimized_mesh['cells']
        
        # Identify poor quality elements
        poor_elements = MeshQuality.identify_poor_quality_elements(
            self.mesh_data, threshold=quality_threshold
        )
        
        # Create a set of vertices to optimize (from poor elements)
        vertices_to_optimize = set()
        for elem_idx in poor_elements:
            vertices_to_optimize.update(cells[elem_idx])
        
        # Build vertex connectivity (neighbors for each vertex)
        vertex_neighbors = [[] for _ in range(len(points))]
        for cell in cells:
            n = len(cell)
            for i in range(n):
                v1 = cell[i]
                v2 = cell[(i + 1) % n]
                if v2 not in vertex_neighbors[v1]:
                    vertex_neighbors[v1].append(v2)
                if v1 not in vertex_neighbors[v2]:
                    vertex_neighbors[v2].append(v1)
        
        # Perform Laplacian smoothing
        for _ in range(iterations):
            # Create a copy of points for this iteration
            new_points = np.copy(points)
            
            # Update each vertex position
            for vertex in vertices_to_optimize:
                neighbors = vertex_neighbors[vertex]
                if neighbors:
                    # Calculate average position of neighbors
                    avg_pos = np.mean([points[n] for n in neighbors], axis=0)
                    # Move vertex towards average position (with relaxation)
                    new_points[vertex] = 0.7 * avg_pos + 0.3 * points[vertex]
            
            # Update points
            points = new_points
            optimized_mesh['points'] = points
        
        # Calculate quality metrics for the optimized mesh
        optimized_quality = MeshQuality.calculate_quality_metrics(optimized_mesh)
        
        # Calculate improvement statistics
        improvement = {
            'original_avg_quality': self.quality_metrics['avg_quality'],
            'optimized_avg_quality': optimized_quality['avg_quality'],
            'improvement_percentage': (
                (self.quality_metrics['avg_quality'] - optimized_quality['avg_quality']) / 
                self.quality_metrics['avg_quality'] * 100
            ) if self.quality_metrics['avg_quality'] > 0 else 0
        }
        
        return {
            'mesh_data': optimized_mesh,
            'quality_metrics': optimized_quality,
            'improvement': improvement
        }
    
    def _smart_smoothing(self, iterations, quality_threshold):
        """
        Optimize mesh using smart smoothing
        
        This method only moves vertices if it improves the quality of adjacent elements.
        """
        # Create a copy of the mesh data to work with
        optimized_mesh = self._copy_mesh_data(self.mesh_data)
        points = optimized_mesh['points']
        cells = optimized_mesh['cells']
        
        # Identify poor quality elements
        poor_elements = MeshQuality.identify_poor_quality_elements(
            self.mesh_data, threshold=quality_threshold
        )
        
        # Create a set of vertices to optimize (from poor elements)
        vertices_to_optimize = set()
        for elem_idx in poor_elements:
            vertices_to_optimize.update(cells[elem_idx])
        
        # Build vertex-to-element connectivity
        vertex_elements = [[] for _ in range(len(points))]
        for i, cell in enumerate(cells):
            for vertex in cell:
                vertex_elements[vertex].append(i)
        
        # Build vertex connectivity (neighbors for each vertex)
        vertex_neighbors = [[] for _ in range(len(points))]
        for cell in cells:
            n = len(cell)
            for i in range(n):
                v1 = cell[i]
                v2 = cell[(i + 1) % n]
                if v2 not in vertex_neighbors[v1]:
                    vertex_neighbors[v1].append(v2)
                if v1 not in vertex_neighbors[v2]:
                    vertex_neighbors[v2].append(v1)
        
        # Perform smart smoothing
        for _ in range(iterations):
            # Create a copy of points for this iteration
            new_points = np.copy(points)
            
            # Update each vertex position
            for vertex in vertices_to_optimize:
                neighbors = vertex_neighbors[vertex]
                if neighbors:
                    # Calculate average position of neighbors
                    avg_pos = np.mean([points[n] for n in neighbors], axis=0)
                    
                    # Store original position
                    original_pos = np.copy(points[vertex])
                    
                    # Temporarily move vertex to new position
                    points[vertex] = 0.7 * avg_pos + 0.3 * original_pos
                    
                    # Calculate quality before and after move
                    affected_elements = vertex_elements[vertex]
                    
                    # Create sub-mesh with only affected elements
                    sub_mesh_before = {
                        'points': np.copy(points),
                        'cells': [cells[i] for i in affected_elements]
                    }
                    sub_mesh_before['points'][vertex] = original_pos
                    
                    sub_mesh_after = {
                        'points': np.copy(points),
                        'cells': [cells[i] for i in affected_elements]
                    }
                    
                    # Calculate quality metrics
                    quality_before = MeshQuality._calculate_manual_quality_metrics(sub_mesh_before)
                    quality_after = MeshQuality._calculate_manual_quality_metrics(sub_mesh_after)
                    
                    # Revert to original position if quality didn't improve
                    if quality_after['avg_quality'] >= quality_before['avg_quality']:
                        points[vertex] = original_pos
                        new_points[vertex] = original_pos
                    else:
                        new_points[vertex] = points[vertex]
            
            # Update points
            points = new_points
            optimized_mesh['points'] = points
        
        # Calculate quality metrics for the optimized mesh
        optimized_quality = MeshQuality.calculate_quality_metrics(optimized_mesh)
        
        # Calculate improvement statistics
        improvement = {
            'original_avg_quality': self.quality_metrics['avg_quality'],
            'optimized_avg_quality': optimized_quality['avg_quality'],
            'improvement_percentage': (
                (self.quality_metrics['avg_quality'] - optimized_quality['avg_quality']) / 
                self.quality_metrics['avg_quality'] * 100
            ) if self.quality_metrics['avg_quality'] > 0 else 0
        }
        
        return {
            'mesh_data': optimized_mesh,
            'quality_metrics': optimized_quality,
            'improvement': improvement
        }
    
    def _ai_optimization(self, iterations, quality_threshold):
        """
        Optimize mesh using AI-based approach
        
        This is a placeholder for AI-based optimization.
        In a real implementation, this would use machine learning models.
        """
        # For now, just use smart smoothing as a fallback
        print("AI optimization not fully implemented. Using smart smoothing instead.")
        return self._smart_smoothing(iterations, quality_threshold)
    
    def reset(self):
        """Reset to original mesh data"""
        self.mesh_data = self._copy_mesh_data(self.original_mesh_data)
        self.quality_metrics = MeshQuality.calculate_quality_metrics(self.mesh_data) 